import sys
import datetime as dt
sys.path.append(r'code/frame') 
from data_read.get_feature import Get_feature_data,  Get_npz_feature
from data_read.get_label import Get_npz_label, Get_label_data
from training.dl_frame.month_rolling.data_split_month import  Rolling_eval_split_month
from training.dl_frame.month_rolling.ic_record import Ic_record
from training.dl_frame.month_rolling.dl_pre_train import Dl_pre_train
# from data_pare.data_fliter import Data_fliter
import operator
import os
# 保证处理矩阵为n*50

class Quant():
    def __init__(self, args) -> None:
        for k, v in args.time_param.items():
            args.time_param[k] = dt.datetime.strptime(v, '%Y-%m')
        self.args = args 

    def get_data(self, Get_feature=Get_feature_data, Get_label=Get_npz_label):
        gf = Get_feature(self.args)
        feature, codes_f, times_f, self.args.feature_list= gf.get_feature()
        gl = Get_label(self.args)
        label, codes_l, times_l = gl.get_label_reg()
        print('检查feature与labe时间索引是否对齐：', operator.eq(times_f.tolist(), times_l.tolist()))
        print('检查feature与labe品种索引是否对齐：', operator.eq(codes_f.tolist(), codes_l.tolist()))
        print(f'使用特征数量：{len(self.args.feature_list)} 具体为： {self.args.feature_list}')
        if  self.args.varieties is not None:
            print(f'使用品种数量：{len(self.args.varieties)} 具体为： {self.args.varieties}')
        self.args.varieties = codes_f
        return feature, label, codes_f, times_f     

    def pre_finetune(self, codes, times, feature, label,  Dl_pre_train=Dl_pre_train): #Data_mean_std=Data_mean_std,
        data_sp = Rolling_eval_split_month(feature, label,  times, self.args)
        ic_re = Ic_record(codes, self.args)
        dl_train = Dl_pre_train(self.args)
        # data_std = Data_mean_std(data_sp.root)
        feature_train, feature_eval,  label_train, label_eval, times_train, times_eval = data_sp.pre_train_split()
        # feature_train, feature_eval = data_std.two_set(feature_train, feature_eval)
        train_pred, train_Y, eval_pred,  eval_Y = dl_train.pre_train_process( feature_train, feature_eval,  label_train, label_eval,  times_train, times_eval, data_sp.pre_root)
        ic_re.ic_record(eval_pred, eval_Y, data_sp.out_month_root, times_eval, 'eval')
        ic_re.ic_record(train_pred, train_Y, data_sp.out_month_root, times_train,  'train')
        for _  in range(0, data_sp.out_month_num, self.args.rolling_step):
            feature_train, feature_eval, feature_out, label_train, label_eval, label_out, times_train, times_eval, times_out = data_sp.fine_tune_split()
            # data_std = Data_mean_std(data_sp.out_month_root)
            # feature_train, feature_eval, feature_out = data_std.three_set(feature_train, feature_eval, feature_out)         
            train_pred, Y_insampl, eval_pred, eval_Y, out_pred,  out_Y = dl_train.finetune_process(feature_train, feature_eval, feature_out, label_train, label_eval, label_out, times_train, times_eval, times_out, data_sp.out_month_root, data_sp.pre_root)
            ic_re.ic_record(out_pred, out_Y, data_sp.out_month_root, times_out,  'test')
            ic_re.ic_record(train_pred, Y_insampl, data_sp.out_month_root, times_train, 'train')
            ic_re.ic_record(eval_pred, eval_Y, data_sp.out_month_root, times_eval,  'eval')
            ic_re.save_ic(self.args.result_root, data_sp.out_month_list)